package org.study.shardingsphere.boot.algorithm;

import com.alibaba.fastjson2.JSON;
import com.google.common.collect.Range;
import lombok.extern.slf4j.Slf4j;
import org.apache.shardingsphere.sharding.api.sharding.standard.*;

import java.time.LocalDateTime;
import java.util.*;

/**
 * 自定义 Standard 分库策略：根据时间按年分库。
 * 参考 https://blog.csdn.net/weixin_44211703/article/details/122684779
 *
 * @author Administrator
 * @date 2022-03-08
 */
@Slf4j
public class StandardShardingAlgorithm4DbDemo implements StandardShardingAlgorithm<LocalDateTime> {
    /***
     * 精确分片算法
     *
     * @param databaseNames
     * @param preciseShardingValue
     * @return java.lang.String
     */
    @Override
    public String doSharding(Collection<String> databaseNames, PreciseShardingValue<LocalDateTime> preciseShardingValue) {
        log.info("databaseNames：{}，preciseShardingValue：{}", JSON.toJSONString(databaseNames), JSON.toJSONString(preciseShardingValue));
        // 分片键对应的值
        LocalDateTime shardingValue = preciseShardingValue.getValue();
        if (shardingValue == null) {
            throw new UnsupportedOperationException("分片键 " + preciseShardingValue.getColumnName() + " 的值不能为空");
        }
        for (String each : databaseNames) {
            //获取年份
            int year = shardingValue.getYear();
            if (each.endsWith(String.valueOf(year))) {
                // 最终要使用的库名
                return each;
            }
        }
        throw new UnsupportedOperationException("没有匹配到可用的数据库节点");
    }

    /**
     * 范围分片算法
     *
     * @param databaseNames
     * @param rangeShardingValue
     * @return java.util.Collection<java.lang.String>
     */
    @Override
    public Collection<String> doSharding(Collection<String> databaseNames, RangeShardingValue<LocalDateTime> rangeShardingValue) {
        log.info("databaseNames：{}，rangeShardingValue：{}", JSON.toJSONString(databaseNames), JSON.toJSONString(rangeShardingValue));
        // 分片键对应的范围
        Range<LocalDateTime> shardingValueRange = rangeShardingValue.getValueRange();
        // 分别获取分片键最值(最小值、最大值)
        LocalDateTime startDateTime = shardingValueRange.lowerEndpoint();
        LocalDateTime endDateTime = shardingValueRange.upperEndpoint();
        // 分别获取最值所属的年份
        int startYear = startDateTime.getYear();
        int endYear = endDateTime.getYear();
        // 分别获取最值所属的月份
        /*int startMonth = startDateTime.getMonthValue();
        int endMonth = endDateTime.getMonthValue();*/
        Set<String> result = new LinkedHashSet<>();
        if (startDateTime.getNano() <= endDateTime.getNano()) {
            for (String dbName : databaseNames) {
                // 获取数据库名称的最后四位，即年份
                int dbYear = Integer.parseInt(dbName.substring(dbName.length() - 4));
                if (dbYear >= startYear && dbYear <= endYear) {
                    result.add(dbName);
                }
            }
        }
        return result;
    }

    @Override
    public String getType() {
        return "SHARDING_DB_DATE2";
    }


    @Override
    public void init(Properties props) {

    }
}
